/*
 * Copyright (c) 2024 endless-sky
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#pragma once

#include <string>
#include <vector>
#include <map>
#include <mtd/mtd-user.h>
#include <libmtd.h>

namespace El {
namespace Upgrade {

/**
 * @brief Flash管理器类，负责MTD设备的读写操作
 */
class MtdManager {
public:
    /**
     * @brief MTD设备信息结构体
     */
    struct MtdDevice {
        int fd = -1;                  // 设备文件描述符
        struct mtd_dev_info dev = {}; // MTD设备信息
        libmtd_t mtd = nullptr;       // libmtd句柄

        // Default constructor
        MtdDevice() = default;
    };

    /**
     * @brief MTD分区映射结构体
     */
    struct MtdPartitionMap {
        std::string name;     // 分区名称
        std::string device;   // MTD设备路径 (/dev/mtdX)
        uint64_t size;        // 分区大小
        uint64_t erasesize;   // 擦除块大小
        uint64_t min_io_size; // 最小I/O大小
        bool writable;        // 是否可写
    };

    MtdManager();
    ~MtdManager();

    /**
     * @brief 初始化MTD设备库
     * @return 是否成功
     */
    bool InitializeMtdLib();

    /**
     * @brief 配置分区信息
     * @param partitions 分区映射列表
     * @return 是否成功
     */
    bool ConfigurePartitions(const std::vector<MtdPartitionMap> &partitions);

    /**
     * @brief 清理MTD资源
     */
    void CleanupMtdResources();

    /**
     * @brief 打开MTD设备
     * @param device 设备路径
     * @param mtd MTD设备结构
     * @return 是否成功
     */
    bool OpenMtdDevice(const std::string &device, MtdDevice &mtd);

    /**
     * @brief 关闭MTD设备
     * @param mtd MTD设备结构
     */
    void CloseMtdDevice(MtdDevice &mtd);

    /**
     * @brief 擦除MTD设备
     * @param mtd MTD设备结构
     * @param offset 偏移地址
     * @param size 擦除大小
     * @return 是否成功
     */
    bool EraseMtdDevice(MtdDevice &mtd, uint64_t offset, uint64_t size);

    /**
     * @brief 擦除MTD设备（带坏块处理）
     * @param mtd MTD设备结构
     * @param offset 偏移地址
     * @param size 擦除大小
     * @return 是否成功
     */
    bool EraseMtdDeviceWithBadBlocks(MtdDevice &mtd, uint64_t offset, uint64_t size);

    /**
     * @brief 流式写入镜像数据
     * @param partitionName 分区名称
     * @param data 数据指针
     * @param length 数据长度
     * @param isLastChunk 是否最后一块数据
     * @param writeOffsets 写入偏移映射表
     * @param writeLogicalOffsets 逻辑写入偏移映射表
     * @param writeStarted 写入开始状态映射表
     * @param writeBuffers 写入缓冲区映射表
     * @return 是否成功
     */
    bool WriteImageStream(const std::string &partitionName,
                          const uint8_t *data,
                          size_t length,
                          bool isLastChunk,
                          std::map<std::string, uint64_t> &writeOffsets,
                          std::map<std::string, uint64_t> &writeLogicalOffsets,
                          std::map<std::string, bool> &writeStarted,
                          std::map<std::string, std::vector<uint8_t>> &writeBuffers);

    /**
     * @brief 获取分区信息
     * @param name 分区名称
     * @return 分区映射指针，失败返回nullptr
     */
    MtdPartitionMap *GetPartition(const std::string &name);

    /**
     * @brief 设置最后一个错误信息
     * @param error 错误信息
     */
    void SetLastError(const std::string &error);

    /**
     * @brief 获取最后一个错误信息
     * @return 错误信息
     */
    std::string GetLastError() const;

private:
    // MTD分区映射
    std::vector<MtdPartitionMap> partitions_;
    libmtd_t mtdLib_;

    // 错误信息
    std::string lastError_;
};

} // namespace Upgrade
} // namespace El
